# 数据集验证，需要加载的包库
import os
from PIL import Image
from matplotlib import pyplot as plt
import xml
import cv2 as cv
import numpy as np
import json
import math
from scipy.spatial import distance
from collections import OrderedDict
import time
import pandas as pds
from torch.utils.data import DataLoader
from torchvision import models
import torch
import torch.nn as nn
from utilour.datasetLoader import MyDataset
import torch.optim as optimer
from utilour.tool import calaccurary,visualize_cam
from utilour.gradcadpp import GradCAM,GradCAMpp
import torchvision
from torchvision.utils import make_grid, save_image
import shutil
from torch.utils.tensorboard import SummaryWriter
from torchsummary import summary
from functools import reduce
from matplotlib import pyplot as plt
import random 

def cleardir(foldirpath):
    if os.path.exists(foldirpath):
        shutil.rmtree(foldirpath)
    os.mkdir(foldirpath)


def loginfo(fpath,lslist,acclist,plb,lb,ims,epoch):
    with open(fpath,'a',encoding="utf-8") as fp:
        for i in range(len(lslist)):
            fp.write("{},{},{},{},{},{}\n".format(str(epoch).replace("\n",""),
            str(lslist[i]).replace("\n",""),
            str(acclist[i]).replace("\n",""),
            str(lb[i]).replace("\n",""),
            str(plb[i]).replace("\n",""),
            str(ims[i]).replace("\n","")))


def printinfo(epoch,ls,als,sig='train'):
    ls=np.array(ls)
    als=np.array(als)
    meanls,maxls,minls=np.mean(ls),np.max(ls),np.min(ls)
    meanas,maxas,minas=np.mean(als),np.max(als),np.min(als)
    print("epoch:{} meanacc:{} minacc:{} maxacc:{} meanloss:{} minloss:{} maxloss:{} sig:{}".format(epoch,
    "%.4f"%meanas,"%.4f"%minas,"%.4f"%maxas,
    "%.4f"%meanls,"%.4f"%minls,"%.4f"%maxls,
    sig
    ))
    return meanas,meanls

def getNoTrain(trstep):
    Nosteps=[]
    for step in  trstep:
        if random.random()>0.5:
            Nosteps.append(step)
    return Nosteps

def splitstep(trainLoader):
    Nosteps=[]
    trainstep=[]
    vailstep=[]
    teststep=[]
    for step,(data,label,imgpath) in  enumerate(trainLoader):
        v=random.random()
        if v<0.9:
            trainstep.append(step)
        elif v<0.95:
            vailstep.append(step)
        else:
            teststep.append(step)
    print(len(trainstep),len(vailstep),len(teststep))
    return trainstep,vailstep,teststep


if __name__ == "__main__":
    zipfile="/home/gis/gisdata/data/jupyterlabhub/gitcode/cardanger/zipfile"
    modeltrain="/home/gis/gisdata/data/jupyterlabhub/gitcode/cardanger/modeltrain"
    localtime = time.localtime(time.time())
    shelltext="zip  -r  {}/modeltrain{}_{}_{}_{}.zip  {}/* ".format(zipfile,localtime.tm_mon,localtime.tm_mday,localtime.tm_hour,localtime.tm_min,modeltrain)
    print(shelltext)
    os.system(shelltext)
    errimgpath="/home/gis/gisdata/data/jupyterlabhub/gitcode/cardanger/dataset/newcardata/errimglist.csv"
    logdir=os.path.join(modeltrain,"log")
    tfrecorddir=os.path.join(modeltrain,"modelRecord")
    modelcpk=os.path.join(modeltrain,"modelcpk")
    cleardir(logdir)
    cleardir(tfrecorddir)
    cleardir(modelcpk)

    writer=SummaryWriter(tfrecorddir) # 记录训练过程中的情况
    # 加载数据
    csvpath="/home/gis/gisdata/data/jupyterlabhub/gitcode/cardanger/dataset/newcardata/train_test_vail.csv" # 文件路径
    traindataset=MyDataset(csvpath,errimgpath,"train",errload=False)
    vaildataset=MyDataset(csvpath,errimgpath,"vail",errload=False)
    testdataset=MyDataset(csvpath,errimgpath,"test",errload=False)
    trainLoader=DataLoader(dataset=traindataset,batch_size=8,num_workers=8,shuffle=False)
    vailLoader=DataLoader(dataset=vaildataset,batch_size=8,num_workers=8)
    testLoader=DataLoader(dataset=testdataset,batch_size=8,num_workers=8)

    # 加载模型--修改模型
    #model=models.resnet50(pretrained=True)
    model=models.densenet121(pretrained=True)
    #fc=nn.Linear(model.fc.in_features,5)
    fc=nn.Linear(model.classifier.in_features,5)
    model.classifier=fc
    #model.fc=fc
    model.cuda()

    #summary(model,input_size=(3,256,256))
    

    tlstep,vlstep,ttstep=0,0,0
    minvlac,minttac=0.6,0.6
    traindecary=0.3
    
    # 设置自动绘制图形
    
    plt.show()
    fig=plt.figure(figsize=(24,6))
    plt.ion()
    stepepoch=[]
    trmeanls=[]
    testmeanls=[]
    vailmeanls=[]
    dtrmeanls=[]
    
    trstep,vlstep,ttstep=splitstep(trainLoader)
    Nosteps=[]
    # 训练模型
    for epoch in range(150):
        print(time.strftime('%Y-%m-%d %H:%M:%S',time.localtime(time.time())))
        lr=0.001 if epoch>=900 else 0.01
        #Optimizer=optimer.SGD(model.parameters(),lr=lr,momentum=0.9)
        #Optimizer=optimer.ASGD(model.parameters(), lr=0.1)
	    #Optimizer=optimer.SGD(model.parameters(),lr=0.001,momentum=0.9)
        Optimizer=optimer.SGD(model.parameters(),lr=0.001,momentum=0.9)
        model.train()
        CRLoss=nn.CrossEntropyLoss()
        #CRLoss=torch.nn.TripletMarginLoss(margin=1.0, p=2.0, eps=1e-06, swap=False, reduction='mean')
        tlls,tlac,tllb,tlplb,tlims=[],[],[],[],[]
        dtlls,dtlac,dtllb,dtlplb,dtlims=[],[],[],[],[]
        vlls,vlac,vllb,vlplb,vlims=[],[],[],[],[]
        ttls,ttac,ttlb,ttplb,ttims=[],[],[],[],[]
        # 一般区域，一半区域进行训练
        gstep=0
        for step,(data,label,imgpath) in  enumerate(trainLoader):
            Optimizer.zero_grad() # 导数归零
            data,label=data.cuda(),label.cuda()
            tlabel=model(data) # 计算结果
            loss=CRLoss(tlabel,label)
            accurary=calaccurary(tlabel.detach(),label.detach())
            if step in Nosteps:
                dtlls.append(loss.detach().cpu().numpy())
                dtlac.append(accurary[0])
                dtlplb.append(tlabel.detach().cpu().numpy()[0])
                dtllb.append(label.detach().cpu().numpy()[0])
                dtlims.append(str(reduce(lambda  x,y: "{} {}".format(x,y),imgpath)))

            if not step in Nosteps and step in trstep:
                loss.backward()
                Optimizer.step()
                tlls.append(loss.detach().cpu().numpy())
                tlac.append(accurary[0])
                tlplb.append(tlabel.detach().cpu().numpy()[0])
                tllb.append(label.detach().cpu().numpy()[0])
                tlims.append(str(reduce(lambda  x,y: "{} {}".format(x,y),imgpath)))
                writer.add_scalar("train loss",loss,gstep)
                writer.add_scalar("train acc",accurary[0],gstep)
                #accurary=calaccurary(tlabel.detach(),label.detach())
            elif step in vlstep:
                vlls.append(loss.detach().cpu().numpy())
                vlac.append(accurary[0])
                vlplb.append(tlabel.detach().cpu().numpy()[0])
                vllb.append(label.detach().cpu().numpy()[0])
                vlims.append(str(reduce(lambda  x,y: "{} {}".format(x,y),imgpath)))
                writer.add_scalar("vail loss",loss,gstep)
                writer.add_scalar("vail acc",accurary[0],gstep)
            elif step in ttstep:
                ttls.append(loss.detach().cpu().numpy())
                ttac.append(accurary[0])
                ttplb.append(tlabel.detach().cpu().numpy()[0])
                ttlb.append(label.detach().cpu().numpy()[0])
                ttims.append(str(reduce(lambda  x,y: "{} {}".format(x,y),imgpath)))
                writer.add_scalar("test loss",loss,gstep)
                writer.add_scalar("test acc",accurary[0],gstep)

            gstep=gstep+1
        torch.save(model,os.path.join(modelcpk,"model_last.pkl"))

        # 输出日志
        loginfo(os.path.join(logdir,"train.csv"),tlls,tlac,tlplb,tllb,tlims,epoch)
        loginfo(os.path.join(logdir,"vail.csv"),vlls,vlac,vlplb,vllb,vlims,epoch)
        loginfo(os.path.join(logdir,"test.csv"),ttls,ttac,ttplb,ttlb,ttims,epoch)
        #loginfo(os.path.join(logdir,"notrain.csv"),dtlls,dtlac,dtlplb,dtllb,dtlims,epoch)
        # 输出训练信息
        tlac,tlls=printinfo(epoch,tlls,tlac,sig="train")
        vlac,vlls=printinfo(epoch,vlls,vlac,sig="vail")
        ttac,ttls=printinfo(epoch,ttls,ttac,sig="test")
        #dtlac,dtlls=printinfo(epoch,dtlls,dtlac,sig="notrain")
        trmeanls.append(tlls)
        vailmeanls.append(vlls)
        testmeanls.append(ttls)
        stepepoch.append(epoch)
        #dtrmeanls.append(dtlls)

        #if tlls<traindecary*3:
        #    #traindecary=traindecary*0.5
        writer.add_scalar('train mean acc',tlac,epoch)
        writer.add_scalar('vail mean acc',vlac,epoch)
        writer.add_scalar('test mean acc',ttac,epoch)
        #writer.add_scalar('Notrain mean acc',dtlac,epoch)
        writer.add_scalar('train mean loss',tlls,epoch)
        writer.add_scalar('vail mean loss',vlls,epoch)
        writer.add_scalar('test mean loss',ttls,epoch)
        #writer.add_scalar('notrain mean loss',dtlls,epoch)
        #if dtlac<0.4 or dtlls<tlls*0.4:
            #Nosteps=getNoTrain(trainLoader)
            #print("Nostep 已经修改")
        if vlac>minvlac:
            minvlac=vlac
            torch.save(model,os.path.join(modelcpk,"model_minvl_{}_{}.pkl".format(epoch,int(minvlac*10000))))
        if ttac>minttac:
            minttac=ttac
            torch.save(model,os.path.join(modelcpk,"model_mintt_{}_{}.pkl".format(epoch,int(minttac*10000))))
        print("current minvailacc:{} mintestacc:{}  traindecary:{}".format(minvlac,minttac,traindecary))
        
        if epoch>5:
        # 绘制图像
            plt.cla()
            plt.axis([-0.5,151,-0.2,2])
            #plt.plot(np.array(stepepoch[3:]), np.array(dtrmeanls[3:]), '-', marker="*",label='trainmean', linewidth=1,color="pink")
            plt.plot(np.array(stepepoch[3:]), np.array(trmeanls[3:]), '-', marker="*",label='trainmean', linewidth=1,color="green")
            plt.plot(np.array(stepepoch[2:-1]), np.array(testmeanls[2:-1]), '-.', marker=".", label='testmean', linewidth=1,color="red")
            plt.plot(np.array(stepepoch[2:-1]), np.array(vailmeanls[2:-1]), '--', marker="+", label='vailmean', linewidth=1,color="blue") 
            plt.legend()
            plt.pause(0.01)


    writer.close()
